package net.cyue.ort.llm.tokenizer;

import net.cyue.util.StreamUtil;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.File;
import java.io.FileInputStream;
import java.io.IOException;
import java.io.InputStream;
import java.util.Arrays;
import java.util.List;

public abstract class Tokenizer {
    protected final Logger logger = LoggerFactory.getLogger(getClass());

    public static class Result {
        public final String prompt;
        public final String[] tokens;
        public final long[] ids;

        public Result(String[] tokens, long[] ids) {
            this(null, tokens, ids);
        }
        public Result(String prompt, String[] tokens, long[] ids) {
            this.prompt = prompt;
            this.tokens = tokens;
            this.ids = ids;
        }

        @Override
        public String toString() {
            StringBuilder sb = new StringBuilder("Result{");

            if (prompt != null) {
                sb.append("\n  prompt='").append(prompt).append("'");
            }

            sb.append("\n  tokens=[");
            if (tokens != null) {
                for (int i = 0; i < tokens.length; i++) {
                    if (i > 0) sb.append(", ");
                    sb.append("'").append(tokens[i]).append("'");
                }
            }
            sb.append("]");

            sb.append("\n  ids=[");
            if (ids != null) {
                for (int i = 0; i < ids.length; i++) {
                    if (i > 0) sb.append(", ");
                    sb.append(ids[i]);
                }
            }
            sb.append("]");

            sb.append("\n}");
            return sb.toString();
        }
    }

    protected Tokenizer() {}
    public Tokenizer(
        String tokenizerFilePath
    ) throws IOException {
        this(new FileInputStream(tokenizerFilePath));
        if (logger.isDebugEnabled()) {
            logger.debug("Tokenizer config loaded from path: {}", tokenizerFilePath);
        }
    }
    public Tokenizer(
        File tokenizerFile
    ) throws IOException {
        this(new FileInputStream(tokenizerFile));
        if (logger.isDebugEnabled()) {
            logger.debug("Tokenizer config loaded from file: {}", tokenizerFile.getAbsolutePath());
        }
    }
    public Tokenizer(
        InputStream tokenizerStream
    ) throws IOException {
        long startTime = System.currentTimeMillis();
        String configContent = StreamUtil.readAsString(tokenizerStream);
        this.loadConfig(configContent);
        if (logger.isDebugEnabled()) {
            logger.debug(
                "Tokenizer config loaded: size={} chars, time={}ms",
                configContent == null ? 0 : configContent.length(),
                System.currentTimeMillis() - startTime
            );
        }
    }

    /**
     * 加载配置
     * @param configContent tokenizer.json的内容。
     */
    protected abstract void loadConfig(String configContent) throws IOException;
    public abstract String[] tokenize(String text);
    public long[] tokens2IDs(List<String> tokenList) {
        return tokens2IDs(tokenList.toArray(new String[0]));
    }
    public abstract long[] tokens2IDs(String[] tokens);

    public abstract Result encode(String text);
    public String decode(List<Long> tokenIDList) {
        return decode(tokenIDList, true);
    }
    public String decode(List<Long> tokenIDList, boolean skipSpecialTokens) {
        if (tokenIDList == null) {
            throw new NullPointerException("Token ID list cannot be null.");
        }
        logDecodeRequest("decode(List)", tokenIDList.size(), skipSpecialTokens);
        return decode(
            tokenIDList.stream().mapToLong(Long::longValue).toArray(),
            skipSpecialTokens
        );
    }
    public String decode(Long[] tokenIds) {
        return decode(tokenIds, true);
    }
    public String decode(Long[] tokenIds, boolean skipSpecialTokens) {
        if (tokenIds == null) {
            throw new NullPointerException("Token IDs cannot be null.");
        }
        logDecodeRequest("decode(Long[])", tokenIds.length, skipSpecialTokens);
        return decode(
            Arrays.stream(tokenIds).mapToLong(Long::longValue).toArray(),
            skipSpecialTokens
        );
    }

    public String decode(Long tokenID) {
        return decode(tokenID, true);
    }
    public String decode(Long tokenID, boolean skipSpecialTokens) {
        if (tokenID == null) {
            throw new NullPointerException("Token ID cannot be null.");
        }
        logDecodeRequest("decode(Long)", 1, skipSpecialTokens);
        return decode((long)tokenID, skipSpecialTokens);
    }

    public String decode(long tokenID) {
        return decode(tokenID, true);
    }
    public String decode(long tokenID, boolean skipSpecialTokens) {
        logDecodeRequest("decode(long)", 1, skipSpecialTokens);
        return decode(new Long[]{tokenID}, skipSpecialTokens);
    }
    public String decode(long[] tokenIDs) {
        return decode(tokenIDs, true);
    }
    public abstract String decode(long[] tokenIDs, boolean skipSpecialTokens);

    protected void logDecodeRequest(String method, int tokenCount, boolean skipSpecialTokens) {
        if (logger.isDebugEnabled()) {
            logger.debug(
                "{} invoked: tokenCount={}, skipSpecialTokens={}",
                method,
                tokenCount,
                skipSpecialTokens
            );
        }
    }
}
